// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/functions/sparksql/aggregates/regr_replacement_aggregate.h>

#include <pollux/exec/simple_aggregate_adapter.h>

namespace kumo::pollux::functions::aggregate::sparksql {
    namespace {
        class RegrReplacementAggregate {
        public:
            using InputType = Row<double>;
            using IntermediateType =
            Row</*n*/ double,
                /*avg*/ double,
                /*m2*/ double>;
            using OutputType = double;

            static bool toIntermediate(
                exec::out_type<Row<double, double, double> > &out,
                exec::arg_type<double> in) {
                out.copy_from(std::make_tuple(1.0, in, 0.0));
                return true;
            }

            struct AccumulatorType {
                double n{0.0};
                double avg{0.0};
                double m2{0.0};

                explicit AccumulatorType(
                    HashStringAllocator * /*allocator*/,
                    RegrReplacementAggregate * /*fn*/) {
                }

                void addInput(
                    HashStringAllocator * /*allocator*/,
                    exec::arg_type<double> data) {
                    n += 1.0;
                    double delta = data - avg;
                    double deltaN = delta / n;
                    avg += deltaN;
                    m2 += delta * (delta - deltaN);
                }

                void combine(
                    HashStringAllocator * /*allocator*/,
                    exec::arg_type<Row<double, double, double> > other) {
                    POLLUX_CHECK(other.at<0>().has_value());
                    POLLUX_CHECK(other.at<1>().has_value());
                    POLLUX_CHECK(other.at<2>().has_value());

                    double otherN = other.at<0>().value();
                    double otherAvg = other.at<1>().value();
                    double otherM2 = other.at<2>().value();

                    double originN = n;
                    n += otherN;
                    double delta = otherAvg - avg;
                    double deltaN = n == 0.0 ? 0.0 : delta / n;
                    avg += deltaN * otherN;
                    m2 += otherM2 + delta * deltaN * originN * otherN;
                }

                bool writeIntermediateResult(exec::out_type<IntermediateType> &out) {
                    out = std::make_tuple(n, avg, m2);
                    return true;
                }

                bool writeFinalResult(exec::out_type<OutputType> &out) {
                    if (n == 0.0) {
                        return false;
                    }
                    out = m2;
                    return true;
                }
            };
        };

        exec::AggregateRegistrationResult registerRegrReplacement(
            const std::string &name,
            bool withCompanionFunctions,
            bool overwrite) {
            std::vector<std::shared_ptr<exec::AggregateFunctionSignature> > signatures{
                exec::AggregateFunctionSignatureBuilder()
                .returnType("double")
                .intermediateType("row(double, double, double)")
                .argumentType("double")
                .build()
            };
            return exec::registerAggregateFunction(
                name,
                std::move(signatures),
                [name](
            core::AggregationNode::Step step,
            const std::vector<TypePtr> &argTypes,
            const TypePtr &resultType,
            const core::QueryConfig & /*config*/)
            -> std::unique_ptr<exec::Aggregate> {
                    POLLUX_CHECK_EQ(
                        argTypes.size(), 1, "{} takes at most one argument", name);
                    return std::make_unique<
                        exec::SimpleAggregateAdapter<RegrReplacementAggregate> >(
                        step, argTypes, resultType);
                },
                withCompanionFunctions,
                overwrite);
        }
    } // namespace

    void registerRegrReplacementAggregate(
        const std::string &prefix,
        bool withCompanionFunctions,
        bool overwrite) {
        registerRegrReplacement(
            prefix + "regr_replacement", withCompanionFunctions, overwrite);
    }
} // namespace kumo::pollux::functions::aggregate::sparksql
